import socket
import threading
import time
import os
import json
import hashlib

class ChatServer:
    def __init__(self):

        # 服务器状态
        self.is_running = False
        self.clients = {}
        self.file_transfers = {}
        # 新增：文件端口和socket
        self.file_server_socket = None
        self.file_port = 12346  # 文件传输端口，可根据需要修改

        # heartbeat
        self.lock = threading.Lock()  # 新增线程锁
        self.HEARTBEAT_INTERVAL = 30  # 心跳检测间隔（秒）
        self.HEARTBEAT_TIMEOUT = 45  # 心跳超时时间（秒）
        self._init_heartbeat_check()

    def _init_heartbeat_check(self):
        """初始化心跳检测线程"""
        self.heartbeat_thread = threading.Thread(
            target=self._check_heartbeats,
            daemon=True
        )
        self.heartbeat_thread.start()

    def _check_heartbeats(self):
        """定期检查客户端心跳"""
        while self.is_running:
            time.sleep(self.HEARTBEAT_INTERVAL)
            current_time = time.time()

            with self.lock:  # 加锁获取客户端副本
                # 创建副本避免遍历时修改字典
                clients_copy = self.clients.copy()

            for client, info in clients_copy.items():
                last_active = info.get('last_active', 0)
                if current_time - last_active > self.HEARTBEAT_TIMEOUT:
                    self._log(f"客户端 {info['nickname']} 心跳超时，断开连接")
                    self._cleanup_client(client)

    def start_server(self, ip, port):
        """启动服务器"""
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.bind((ip, port))
            self.server_socket.listen(5)
            self.is_running = True
            threading.Thread(target=self._accept_clients, daemon=True).start()
            self._log(f"聊天服务端口启动成功 [{ip}:{port}]")
            # 新增：启动文件传输服务器
            self.start_file_server(ip, self.file_port)
        except Exception as e:
            self._log(f"服务器启动失败: {str(e)}", "ERROR")

    def start_file_server(self, ip, port):
        """启动文件传输服务器"""
        try:
            self.file_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.file_server_socket.bind((ip, port))
            self.file_server_socket.listen(2)
            self._log(f"文件传输端口启动成功 [{ip}:{port}]")
            threading.Thread(target=self._accept_file_clients, daemon=True).start()
        except Exception as e:
            self._log(f"文件传输端口启动失败: {str(e)}", "ERROR")

    def stop_server(self):
        """停止服务器"""
        self.is_running = False
        try:
            for client in self.clients.copy():
                client.close()
            self.server_socket.close()
            # 新增：关闭文件传输socket
            if self.file_server_socket:
                self.file_server_socket.close()
        except:
            pass
        self._log("服务器已停止")

    def _accept_clients(self):
        """接受客户端连接"""
        while self.is_running:
            try:
                client_socket, addr = self.server_socket.accept()
                threading.Thread(
                    target=self._handle_client,
                    args=(client_socket,),
                    daemon=True
                ).start()
                self._log(f"新连接来自: [{addr[0]}:{addr[1]}]")
            except:
                break

    def _accept_file_clients(self):
        """接受文件客户端连接"""
        while self.is_running:
            if self.file_server_socket is None:
                break
            try:
                file_client_socket, addr = self.file_server_socket.accept()
                threading.Thread(
                    target=self._handle_file_client,
                    args=(file_client_socket,),
                    daemon=True
                ).start()
                self._log(f"新文件连接来自: [{addr[0]}:{addr[1]}]")
            except:
                break

    def _broadcast(self, message, sender_socket=None):
        """广播消息到所有在线主端口客户端"""
        for client in self.clients.copy():
            if sender_socket is None or client != sender_socket:
                try:
                    client.send(message)
                except:
                    self._cleanup_client(client)

    def _handle_client(self, client_socket):
        """处理主端口客户端连接（仅聊天和心跳）"""
        try:
            client_ip, client_port = client_socket.getpeername()
            nickname = client_socket.recv(1024).decode().strip()
            client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096)
            with self.lock:
                self.clients[client_socket] = {
                    "nickname": nickname,
                    "ip": f"{client_ip}:{client_port}",
                    "last_active": time.time()
                }
                self._update_user_list()
                self._log(f"用户上线: '{nickname}'  [{client_ip}:{client_port}]", "Login")
            while self.is_running:
                try:
                    data = client_socket.recv(4096)
                    if not data:
                        self._log(f"客户端 {nickname} 断开连接", "ERROR")
                        break
                    with self.lock:
                        if client_socket in self.clients:
                            self.clients[client_socket]["last_active"] = time.time()
                    if data == b"HEARTBEAT_REQUEST":
                        client_socket.send(b"HEARTBEAT_RESPONSE")
                        continue
                    # 聊天消息处理
                    try:
                        message = data.decode('utf-8')
                        self._broadcast(data, client_socket)
                        self._log(f"[{nickname}] {message}")
                    except UnicodeDecodeError:
                        self._log(f"收到无效消息格式", "ERROR")
                except socket.timeout:
                    continue
                except Exception as e:
                    self._log(f"处理消息时发生错误: {str(e)}", "ERROR")
                    break
        except Exception as e:
            self._log(f"客户端连接处理错误: {str(e)}", "ERROR")
        finally:
            self._cleanup_client(client_socket)

    def _handle_file_client(self, file_client_socket):
        """处理文件客户端连接（接收文件或支持下载）"""
        try:
            # 先接收一行，判断是上传还是下载
            header_buffer = b""
            while b"\n" not in header_buffer:
                chunk = file_client_socket.recv(1024)
                if not chunk:
                    self._log("文件客户端连接断开（未收到完整请求）", "ERROR")
                    return
                header_buffer += chunk
            header_line, remain = header_buffer.split(b"\n", 1)
            header = header_line.decode('utf-8')

            if header.startswith("FILE_INIT:"):
                # 上传流程（与之前一致）
                _, file_name, file_size, md5_hex, file_sender = header.split(":", 4)
                file_name = file_name.strip()
                file_size = int(file_size)
                md5_hex = md5_hex.strip()
                file_sender = file_sender.strip()
                if file_size <= 0 or file_size > 1024 * 1024 * 100:
                    self._log(f"无效的文件大小: {file_size} bytes", "ERROR")
                    return
                self._log(f"开始接收文件: {file_name} ({file_size} bytes) (MD5: {md5_hex})")
                received = len(remain)
                file_data = bytearray(remain)
                while received < file_size:
                    chunk = file_client_socket.recv(min(4096, file_size - received))
                    if not chunk:
                        self._log(f"文件接收中断: {file_name}", "ERROR")
                        return
                    file_data.extend(chunk)
                    received += len(chunk)
                    if received % (1024*1024) == 0 or received == file_size:
                        progress = received / file_size * 100
                        self._log(f"接收 {file_name}: {progress:.1f}% ({received/1024:.1f}KB/{file_size/1024:.1f}KB)")
                received_md5 = hashlib.md5(file_data).hexdigest()
                if received_md5 != md5_hex:
                    self._log(f"文件校验失败: {file_name} (预期:{md5_hex} 实际:{received_md5})", "ERROR")
                    return
                save_dir = "server_files"
                os.makedirs(save_dir, exist_ok=True)
                save_path = os.path.join(save_dir, file_name)
                if os.path.exists(save_path):
                    name, ext = os.path.splitext(file_name)
                    timestamp = time.strftime("%Y%m%d_%H%M%S")
                    file_name = f"{name}_{timestamp}{ext}"
                    save_path = os.path.join(save_dir, file_name)
                with open(save_path, "wb") as f:
                    f.write(file_data)
                self._log(f"文件保存成功: {save_path} (MD5: {received_md5})")
                try:
                    client_ip, client_port = file_client_socket.getpeername()
                    sender_ip = f"{client_ip}:{client_port}"
                except:
                    sender_ip = "未知"
                notice = json.dumps({
                    "type": "file_notice",
                    "file_name": file_name,
                    "file_size": file_size,
                    "md5": md5_hex,
                    "sender": file_sender,
                    "ip": sender_ip
                }).encode('utf-8')
                self._broadcast(notice, None)
            elif header.startswith("GET:"):
                # 下载流程
                file_name = header[4:].strip()
                save_dir = "server_files"
                file_path = os.path.join(save_dir, file_name)
                if not os.path.exists(file_path):
                    file_client_socket.sendall(b"ERROR:FileNotFound\n")
                    self._log(f"客户端下载文件不存在: {file_name}", "ERROR")
                    return
                file_size = os.path.getsize(file_path)
                with open(file_path, "rb") as f:
                    file_data = f.read()
                md5_hex = hashlib.md5(file_data).hexdigest()
                # 发送文件头
                header = f"FILE_INIT:{file_name}:{file_size}:{md5_hex}\n".encode('utf-8')
                file_client_socket.sendall(header)
                # 发送文件内容
                sent = 0
                while sent < file_size:
                    chunk = file_data[sent:sent+4096]
                    file_client_socket.sendall(chunk)
                    sent += len(chunk)
                self._log(f"客户端下载完成: {file_name} ({file_size} bytes)")
            else:
                self._log(f"收到无效文件端口请求: {header}", "ERROR")
        except Exception as e:
            self._log(f"文件客户端处理错误: {str(e)}", "ERROR")
        finally:
            try:
                file_client_socket.close()
            except:
                pass

    def _cleanup_client(self, client_socket):
        """清理客户端资源"""
        if client_socket in self.clients:
            info = self.clients[client_socket]
            try:
                # 获取客户端地址信息（备选方案）
                client_ip, client_port = client_socket.getpeername()
                address = f"{client_ip}:{client_port}"
            except:
                # 如果socket已关闭，使用存储的地址
                address = info.get('ip', '未知地址')
            del self.clients[client_socket]
            self._update_user_list()
            self._log(f"用户离线: '{info['nickname']}'  [{address}]", "ERROR")
        try:
            client_socket.close()
        except:
            pass

    def _update_user_list(self):
        """更新在线用户列表"""
        users = [f"{info['nickname']} \n[{info['ip']}]\n" for info in self.clients.values()]

    def _log(self, message, level="INFO"):
        """控制台日志输出"""
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        color_code = ""
        if level == "ERROR":
            color_code = "\033[91m"  # 红色
        elif level == "WARNING":
            color_code = "\033[93m"  # 黄色
        elif level == "Login":
            color_code = "\033[92m"  # 绿色

        print(f"{color_code}[{timestamp}] {message}\033[0m")

    def run_cli(self):
        """控制台交互界面"""
        print("""Chat Server 控制台
        可用命令：
        /start    - 启动服务器
        /stop     - 停止服务器
        /status   - 查看状态
        /users    - 查看在线用户
        /clear    - 清空日志
        /exit     - 退出程序
        """)

        while True:
            cmd = input("> ").strip().lower()
            if cmd == "/start":
                self.start_server("192.168.1.16", 12345)
            elif cmd == "/stop":
                self.stop_server()
            elif cmd == "/status":
                print(f"运行状态: {'运行中' if self.is_running else '已停止'}")
            elif cmd == "/users":
                print(f"在线用户 ({len(self.clients)})：")
                for info in self.clients.values():
                    print(f"- {info['nickname']} [{info['ip']}]")
            elif cmd == "/clear":
                os.system('cls' if os.name == 'nt' else 'clear')
            elif cmd == "/exit":
                self.stop_server()
                break


if __name__ == "__main__":
    server = ChatServer()
    server.run_cli()
